#!/usr/bin/env python3
# -*- encoding: utf-8 -*-
import numpy as np
from ai_library.components.findPlate import detect
from ai_library.components.splitPlate import *
from ai_library.components.svm_predict import test

class PlateRec():
    def __init__(self):
        pass

    def picturePreprocessing(self, roi):
        # 车牌图像预处理
        test = detect(roi)
        if test is None or test.shape[0] < 20:
            return

        gray = cv.cvtColor(test, cv.COLOR_BGR2GRAY)
        ret, binary = cv.threshold(gray, 0, 255, cv.THRESH_BINARY + cv.THRESH_OTSU)
        upImg = cutMatUp(binary, 0)
        downImg = cutMatDown(upImg, 0)
        leftImg = cutMapLeft(downImg, 0)
        rightImg = cutMatRight(leftImg, 0)
        h, w = rightImg.shape[:2]
        temp = rightImg[int(0.020 * h):int(h * 0.97), int(0.020 * w): int(0.98 * w)]
        clearLiuDing(temp)
        upImg = cutMatUp(temp, 255)
        downImg = cutMatDown(upImg, 255)
        leftImg = cutMapLeft(downImg, 255)
        rightImg = cutMatRight(leftImg, 255)
        resize = cv.resize(rightImg, (144, 34), interpolation = cv.INTER_AREA)
        return resize

    def findelPlate(self, frame):
        # 查找车牌
        blue_lower = np.array([100, 100, 80])
        blue_upper = np.array([124, 255, 255])
        kernel = np.ones((5, 5), np.float32)
        hsv = cv.cvtColor(frame, cv.COLOR_BGR2HSV)
        blue_mask = cv.inRange(hsv, blue_lower, blue_upper)
        mask_blue_dilate = cv.dilate(blue_mask, kernel, iterations = 3)
        res = cv.bitwise_and(frame, frame, mask = mask_blue_dilate)
        gray = cv.cvtColor(res, cv.COLOR_BGR2GRAY)
        _, threshed = cv.threshold(gray, 0, 255, cv.THRESH_BINARY + cv.THRESH_OTSU)
        contours, hierarchy = cv.findContours(threshed, cv.RETR_TREE, cv.CHAIN_APPROX_SIMPLE)[-2:]
        cnt_t = None
        area_max = 0
        for cnt in contours:
            Area = cv.contourArea(cnt)
            if Area > area_max:
                area_max = Area
                cnt_t = cnt
        if cnt_t is None and area_max < 800:
            return

        x, y, w, h = cv.boundingRect(cnt_t)
        roi = frame.copy()[y:y + h, x:x + w]
        return roi

    def plateRecognition(self, frame):
        # 车牌识别程序
        plate_roi = self.findelPlate(frame)
        if isinstance(plate_roi, np.ndarray):
            resize = self.picturePreprocessing(plate_roi)
            if isinstance(resize, np.ndarray):
                plate_str = test(resize)
                return plate_str

import os
def main():
    imagePath = './download'
    plate_rec = PlateRec()

    for img_path in os.listdir(imagePath):
        img_path = os.path.join(imagePath, img_path)

        font = cv.FONT_HERSHEY_SIMPLEX
        img_plate = cv.imread(img_path)
        plate_str = plate_rec.plateRecognition(img_plate)
        cv.putText(img_plate, plate_str, (20, 200), font, .8, (255, 0, 255), 1)
        cv.imshow('plate', img_plate)
        cv.waitKey(0)

if __name__ == '__main__':
    main()